import pickle
import pymongo
import time
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA


# 创建mongodb连接
def get_mongo_connection(col_name):
    client = pymongo.MongoClient("129.204.20.87", 27017)
    db = client["questionnaire"]
    db.authenticate('admin', '123456')
    col = db[col_name]
    return col


# 从mongodb获取数据
def get_data():
    collection = get_mongo_connection("questionnaire")

    res_list = []
    for i in range(1, 1001):
        print(f"handling the {i}")
        time.sleep(0.5)
        query = {"_id": {"$eq": f"{i}"}}
        for x in collection.find(query):
            res_list.append(x)

    with open("./data.obj", "wb") as f:
        pickle.dump(res_list, f)


# 提炼数据
def extract_data():
    X_data = []
    with open("./data.obj", "rb") as file:
        questionnaire_list = pickle.load(file)

        for questionnaire in questionnaire_list:
            X = []
            for question in questionnaire["questionList"]:
                for choice in question["choices"].items():
                    X.append(choice[1])

            X_data.append(X)

    X_data = np.array(X_data)
    return X_data


# 根据性别划分数据
def gender_division(X):
    male = []
    female = []
    for x in X:
        if x[0] == 1:
            male.append(x)
        else:
            female.append(x)

    return np.array(male), np.array(female)


# K-means 算法训练
def k_means_training(_X):
    # 降维
    pca = PCA(n_components=2)
    X_2D = pca.fit_transform(_X)
    km = KMeans(n_clusters=12, max_iter=100000)
    y_pre = km.fit_predict(X_2D)

    plt.scatter(X_2D[:, 0], X_2D[:, 1], c=y_pre)
    plt.show()

    return y_pre


# 进行宿舍匹配
def divide_data_into_dormitory(_X, male_data, female_data):
    male = {}
    female = {}
    male_count = 0
    female_count = 0
    for user_id, question in enumerate(_X):
        gender = question[0]
        if gender == 1:
            male[f"{user_id + 1}"] = male_data[male_count]
            male_count += 1
        else:
            female[f"{user_id + 1}"] = female_data[female_count]
            female_count += 1

    male_dormitory = {}
    for man_key in male.keys():
        if male_dormitory.get(male[man_key]) is None:
            male_dormitory[male[man_key]] = []
            male_dormitory[male[man_key]].append(man_key)
        else:
            male_dormitory[male[man_key]].append(man_key)

    female_dormitory = {}
    for woman_key in female.keys():
        if female_dormitory.get(female[woman_key]) is None:
            female_dormitory[female[woman_key]] = []
            female_dormitory[female[woman_key]].append(woman_key)
        else:
            female_dormitory[female[woman_key]].append(woman_key)

    male_dormitory_result = {}
    for key in male_dormitory:
        for i, j in enumerate(male_dormitory[key]):
            if (i % 4) == 0:
                male_dormitory_result[j] = male_dormitory[key][i + 1: i + 4]
            elif (i % 4) == 1:
                male_dormitory_result[j] = male_dormitory[key][i - 1: i] + male_dormitory[key][i + 1: i + 3]
            elif (i % 4) == 2:
                male_dormitory_result[j] = male_dormitory[key][i - 2: i] + male_dormitory[key][i + 1: i + 2]
            else:
                male_dormitory_result[j] = male_dormitory[key][i - 3: i]

    female_dormitory_result = {}
    for key in female_dormitory:
        for i, j in enumerate(female_dormitory[key]):
            if (i % 4) == 0:
                female_dormitory_result[j] = female_dormitory[key][i + 1: i + 4]
            elif (i % 4) == 1:
                female_dormitory_result[j] = female_dormitory[key][i - 1: i] + female_dormitory[key][i + 1: i + 3]
            elif (i % 4) == 2:
                female_dormitory_result[j] = female_dormitory[key][i - 2: i] + female_dormitory[key][i + 1: i + 2]
            else:
                female_dormitory_result[j] = female_dormitory[key][i - 3: i]

    dormitory_result = {**male_dormitory_result, **female_dormitory_result}
    res = []
    for key in dormitory_result.keys():
        dorm = {"userId": key, "roomMate": dormitory_result[key]}
        res.append(dorm)
    return res


# 将数据传到mongodb
def post_data(data):
    collection = get_mongo_connection("dormitory")
    collection.insert_many(result)


if __name__ == '__main__':
    # get_data()
    X = extract_data()
    male_X, female_X = gender_division(X)
    male_predict = k_means_training(male_X)
    female_predict = k_means_training(female_X)
    result = divide_data_into_dormitory(X, male_predict, female_predict)
    print("success")